using System;
using System.IO;
using System.ServiceModel;
using System.ServiceModel.Channels;
using System.ServiceModel.Dispatcher;

namespace Universe.WCF.Behaviors
{
    class TrafficInspector : IDispatchMessageInspector, IClientMessageInspector
    {
        private TrafficStatisticConfiguration _configuration;

        public TrafficInspector(TrafficStatisticConfiguration configuration)
        {
            _configuration = configuration;
        }

        private static Message GetMessageSize(MessageBuffer buffer, out int size)
        {
            // Message msg = buffer.CreateMessage();
            // Console.WriteLine("\n{0}\n", msg);
            using (MemoryStream mem = new MemoryStream())
            {
                buffer.WriteMessage(mem);
                // Console.WriteLine("RAW" + Environment.NewLine + Encoding.ASCII.GetString(mem.ToArray()));
                size = (int) mem.Length;
                return buffer.CreateMessage();
            }
        }

        object IDispatchMessageInspector.AfterReceiveRequest(ref Message request,
                                                             IClientChannel channel,
                                                             InstanceContext instanceContext)
        {
            int size;
            request = GetMessageSize(request.CreateBufferedCopy(int.MaxValue), out size);
            string action = request.Headers.Action;
            _configuration.OnServer(action, 1, size, 0);
            return action;
        }

        void IDispatchMessageInspector.BeforeSendReply(ref Message reply, object correlationState)
        {
            int size;
            reply = GetMessageSize(reply.CreateBufferedCopy(int.MaxValue), out size);
            string action = correlationState as string;
            if (action == null)
                action = reply.Headers.Action;

            _configuration.OnServer(action, 0, 0, size);
        }

        void IClientMessageInspector.AfterReceiveReply(ref Message reply, object correlationState)
        {
            int size;
            string action = correlationState as string;
            if (action == null)
                action = reply.Headers.Action;

            reply = GetMessageSize(reply.CreateBufferedCopy(int.MaxValue), out size);
            _configuration.OnClient(action, 0, size, 0);
        }

        object IClientMessageInspector.BeforeSendRequest(ref Message request,
                                                         IClientChannel channel)
        {
            int size;
            request = GetMessageSize(request.CreateBufferedCopy(int.MaxValue), out size);
            string action = request.Headers.Action;
            _configuration.OnClient(action, 1, 0, size);
            return action;
        }
    }
}